import io
import pathlib

import PIL.Image
import matplotlib.animation
import matplotlib.pyplot

import main_infer
import shared.checkpoint


def main():
    checkpoint_manager = shared.checkpoint.CheckpointManager(pathlib.Path("./checkpoints"))
    checkpoints = checkpoint_manager.existed()

    images = []
    for key in sorted(checkpoints):
        print(f"Generating image for checkpoint {key}...")

        checkpoint = checkpoints[key]
        stream = io.BytesIO()

        figure = main_infer.plot_checkpoint(checkpoint)
        figure.suptitle(f"Epoch {key:0000000}", x="0.9")
        figure.savefig(stream)
        matplotlib.pyplot.close(figure)

        stream.seek(0)
        image = PIL.Image.open(stream)
        images.append(image)

    print(f"Generating animation...")
    first_image = images.pop(0)
    first_image.save("./checkpoints/animation.gif",
                     format="gif",
                     save_all=True,
                     append_images=images)


if __name__ == "__main__":
    main()
